from .Dataset import Dataset
from ..utils.dataset_utils import parse_im_name
from collections import defaultdict
import numpy as np
from  PIL import Image
import os.path as osp
class TrainSet(Dataset):
    def __init__(self,im_dir=None,im_names=None,ids2labels=None,ids_per_batch=None,ims_per_id=None,**kwargs):
        self.im_dir = im_dir
        self.im_names = im_names
        self.ids2labels = ids2labels
        self.ims_per_id = ims_per_id
        im_ids = [parse_im_name(name, 'id') for name in im_names]
        self.ids_to_im_inds = defaultdict(list)
        for ind, id in enumerate(im_ids):
            self.ids_to_im_inds[id].append(ind)
        self.ids = list(self.ids_to_im_inds.keys())# in python2: self.ids[ptr], while in python3: list(self.ids)[ptr]
        super(TrainSet,self).__init__(dataset_size=len(self.ids),batch_size=ids_per_batch,**kwargs)

    def next_batch(self):
        '''Next batch of images and labels.
        :return:
        ims: numpy array with shape[N,C,W,H]
        img_names: a numpy array of image names, len(img_names)>=1
        labels: a numpy array of image labels, len(labels)>=1
        mirrored: a numpy array of booleans, whether the images are mirrored
        self.eopch_done: whether the epoch is over
        '''
        # Start enqueuing and other preparation at the beginning of an epoch.
        if self.epoch_done and self.shuffle:
            np.random.shuffle(self.ids)
        samples, self.epoch_done = self.prefetcher.next_batch()
        im_list, im_names, labels, mirrored = zip(*samples)
        # Transform the list into a numpy array with shape[N,C,H,W]
        ims = np.stack(np.concatenate(im_list))
        im_names = np.concatenate(im_names)
        labels = np.concatenate(labels)
        mirrored = np.concatenate(mirrored)
        return ims, im_names, labels, mirrored, self.epoch_done

    def get_sample(self, ptr):
        """
        Here one sample means several images (and labels etc) of one id.
        :return:
        ims: a list of images
        """
        inds = self.ids_to_im_inds[self.ids[ptr]]
        if len(inds) < self.ims_per_id:
            inds = np.random.choice(inds, self.ims_per_id, replace=True)
        else:
            inds = np.random.choice(inds, self.ims_per_id, replace=False)
        im_names = [self.im_names[ind] for ind in inds]
        ims = [np.asarray(Image.open(osp.join(self.im_dir, name))) for name in im_names]
        ims, mirrored = zip(*[self.pre_process_im(im) for im in ims])
        # print("ids2labels",self.ids2labels)
        labels = [self.ids2labels[self.ids[ptr]] for _ in range(self.ims_per_id)]
        # print("lds2label:{}".format(self.ids2labels[self.ids[ptr]]))
        return  ims, im_names, labels, mirrored
